import time

import torch
from tianshou.data import Batch
from tianshou.policy import DQNPolicy
from tianshou.utils.net.common import Net

from baseline.tianshou.test_dqn import get_args
from ship_envs.environment.env_factory import get_single_env_by_task


def show():
    args = get_args()
    env = get_single_env_by_task(args.task)
    args.state_shape = env.observation_space.shape or env.observation_space.n
    args.action_shape = env.action_space.shape or env.action_space.n
    args.resume_path = './log/case_3/dqn/policy.pth'
    net = Net(
        args.state_shape,
        args.action_shape,
        hidden_sizes=args.hidden_sizes,
        device=args.device,
    ).to(args.device)
    optim = torch.optim.Adam(net.parameters(), lr=args.lr)
    policy = DQNPolicy(
        net,
        optim,
        args.gamma,
        args.n_step,
        target_update_freq=args.target_update_freq
    )

    # load a previous policy
    if args.resume_path:
        policy.load_state_dict(torch.load(args.resume_path, map_location=args.device))


    obs = env.reset()
    flag = True
    while True:
        # if keyboard.read_key() == "space":
        #     print("You pressed space")
        #     flag = not flag
        if flag:
            batch = Batch(obs=[obs], info={})
            res = policy(batch)
            # if res.act[0] == 1:
            #     print('turn right')
            # elif res.act[0] == 2:
            #     print('turn left')
            obs, reward, done, info = env.step(res.act[0])
            print(obs, reward)
            env.render()
            time.sleep(0.1)
            if done:
                break





if __name__ == '__main__':
    show()